"""
    重写 from rest_framework.viewsets import ModelViewSet 的 返回格式
    或者直接到 ModelViewSet 中 改写源码，只需更改 return
"""

from rest_framework import status
from rest_framework import viewsets
from rest_framework.decorators import action
from .json_response import JsonResponse
from utils.export import ExportExcel
from audit.views import operation


class CustomModelViewSet(viewsets.ModelViewSet):

    sheet = ''  # 文件名
    title = ''  # 表头

    def create(self, request, *args, **kwargs):
        """
        创建一个模型实例.
        """
        serializer = self.get_serializer(data=request.data)
        serializer.is_valid(raise_exception=True)
        self.perform_create(serializer)
        operation.operation(request, serializer.data['name'], serializer.instance._meta.verbose_name_plural, '创建')
        headers = self.get_success_headers(serializer.data)
        return JsonResponse(data=serializer.data, code=200, message="成功", status=status.HTTP_201_CREATED,
                            headers=headers)

    def list(self, request, *args, **kwargs):
        """
        列出 queryset.
        """
        queryset = self.filter_queryset(self.get_queryset())
        page = self.paginate_queryset(queryset)
        if page is not None:
            serializer = self.get_serializer(page, many=True)
            return self.get_paginated_response(serializer.data)

        serializer = self.get_serializer(queryset, many=True)
        return JsonResponse(data=serializer.data, code=200, message="成功", status=status.HTTP_200_OK)

    def retrieve(self, request, *args, **kwargs):
        """
        检索模型实例.
        """
        instance = self.get_object()
        serializer = self.get_serializer(instance)
        return JsonResponse(data=serializer.data, code=200, message="成功", status=status.HTTP_200_OK)

    def update(self, request, *args, **kwargs):
        """
        更新模型实例.
        """
        partial = kwargs.pop('partial', False)
        instance = self.get_object()
        serializer = self.get_serializer(instance, data=request.data, partial=partial)
        serializer.is_valid(raise_exception=True)
        self.perform_update(serializer)

        if getattr(instance, '_prefetched_objects_cache', None):
            # 如果'prefetch_related'已经应用到queryset，需要强制取消实例上的预取缓存
            instance._prefetched_objects_cache = {}
        operation.operation(request, instance, instance._meta.verbose_name_plural, '更新')
        return JsonResponse(data=serializer.data, code=200, message="成功", status=status.HTTP_200_OK)

    def destroy(self, request, *args, **kwargs):
        """
        删除一个模型实例.
        """
        instance = self.get_object()
        self.perform_destroy(instance)
        operation.operation(request, instance, instance._meta.verbose_name_plural, '删除')
        return JsonResponse(data=[], code=200, message="成功", status=status.HTTP_200_OK)

    @action(methods=['post'], detail=False)  # url 访问127.0.0.1:8080/api/方法名
    def export_all(self, request):
        """
        根据数据表导出Excel表
        """
        queryset = self.filter_queryset(self.get_queryset())
        serializer = self.get_serializer(queryset, many=True)
        name = ExportExcel(serializer.data, self.sheet, self.title).write(serializer.data, self.sheet, self.title)
        operation.operation(request, name, self.get_queryset().model._meta.verbose_name_plural, '导出所有')
        return JsonResponse(data={'name': name}, code=200, message='  Excel 已导出')

    @action(methods=['post'], detail=False)  # url 访问127.0.0.1:8080/api/方法名
    def export_current(self, request):
        """
        根据前端传来的搜索内容、页码等条件导出Excel表
        """
        queryset = self.filter_queryset(self.get_queryset())
        page = self.paginate_queryset(queryset)
        if page is not None:
            serializer = self.get_serializer(page, many=True)
            name = ExportExcel(serializer.data, self.sheet, self.title).write(serializer.data, self.sheet, self.title)
            operation.operation(request, name, self.get_queryset().model._meta.verbose_name_plural, '导出部分')
            return JsonResponse(data={'name': name}, code=200, message='  Excel 已导出')
        else:
            return JsonResponse(data={}, code=400, message='  Excel 导出错误')

